"""
机器学习专用插件
"""
from typing import Any

# 修改导入语句，使用绝对导入而不是相对导入
from plugins.base_plugin import BasePlugin
import time
import random
import json
from core.logger import get_logger

# 获取日志记录器
logger = get_logger(__name__)


class DataLoader(BasePlugin):
    """数据加载插件"""

    def execute(self, input_data: Any) -> dict:
        """模拟数据加载"""
        node_config = self.config.get("config", {})
        data_path = node_config.get("data_path", "default_data.csv")
        sample_size = node_config.get("sample_size", 1000)

        # 生成模拟数据
        data = []
        # 从data_path的csv文件中读取数据
        import pandas as pd
        df = pd.read_csv(data_path)
        data = df.to_dict(orient="records")

        return {
            "raw_data": data,
            "data_source": data_path,
            "sample_size": sample_size,
            "loaded_by": self.plugin_id
        }


class FeatureEngineer(BasePlugin):
    """特征工程插件"""

    def execute(self, input_data: Any) -> dict:
        """执行特征工程"""
        if isinstance(input_data, dict) and "raw_data" in input_data:
            raw_data = input_data["raw_data"]
        else:
            raw_data = input_data

        feature_engineering_steps = self.config.get("steps", ["normalize", "encode"])
        time.sleep(0.15)  # 模拟处理时间

        # 模拟特征工程处理
        processed_features = []
        for record in raw_data:
            processed_record = record.copy()
            # 数值特征标准化
            if "normalize" in feature_engineering_steps:
                processed_record["feature1_normalized"] = record["feature1"] / 100.0
                processed_record["feature2_normalized"] = record["feature2"] / 50.0

            # 类别特征编码
            if "encode" in feature_engineering_steps:
                processed_record["feature3_encoded"] = {"A": 0, "B": 1, "C": 2}[record["feature3"]]

            processed_features.append(processed_record)

        feature_metadata = {
            "original_features": ["feature1", "feature2", "feature3"],
            "engineered_features": ["feature1_normalized", "feature2_normalized", "feature3_encoded"],
            "steps_applied": feature_engineering_steps,
            "processed_count": len(processed_features)
        }

        return {
            "processed_data": processed_features,
            "feature_metadata": feature_metadata,
            "engineered_by": self.plugin_id
        }


class TrainTestSplitter(BasePlugin):
    """训练测试数据分割插件"""

    def execute(self, input_data: Any) -> dict:
        """分割训练和测试数据"""
        if isinstance(input_data, dict) and "processed_data" in input_data:
            data = input_data["processed_data"]
        else:
            data = input_data

        test_size = self.config.get("test_size", 0.2)
        random_state = self.config.get("random_state", 42)
        time.sleep(0.05)

        # 模拟数据分割
        split_index = int(len(data) * (1 - test_size))
        train_data = data[:split_index]
        test_data = data[split_index:]

        # 分离特征和标签
        X_train = [{k: v for k, v in item.items() if k != "target"} for item in train_data]
        y_train = [item["target"] for item in train_data]
        X_test = [{k: v for k, v in item.items() if k != "target"} for item in test_data]
        y_test = [item["target"] for item in test_data]

        split_info = {
            "train_size": len(train_data),
            "test_size": len(test_data),
            "test_ratio": test_size,
            "random_state": random_state
        }

        return {
            "X_train": X_train,
            "y_train": y_train,
            "X_test": X_test,
            "y_test": y_test,
            "split_info": split_info,
            "split_by": self.plugin_id
        }


class ModelTrainer(BasePlugin):
    """模型训练插件"""

    def execute(self, input_data: Any) -> dict:
        """训练机器学习模型"""
        X_train = input_data["X_train"]
        y_train = input_data["y_train"]

        model_type = self.config.get("model_type", "logistic_regression")
        hyperparameters = self.config.get("hyperparameters", {})
        time.sleep(0.2)  # 模拟训练时间

        # 模拟模型训练
        training_metrics = {
            "train_accuracy": random.uniform(0.85, 0.95),
            "train_f1_score": random.uniform(0.80, 0.90),
            "training_time": 0.2
        }

        # 模拟训练好的模型（实际项目中会是模型对象）
        trained_model = {
            "model_type": model_type,
            "feature_count": len(X_train[0]) if X_train else 0,
            "sample_count": len(X_train),
            "hyperparameters": hyperparameters
        }

        model_metadata = {
            "model_type": model_type,
            "hyperparameters": hyperparameters,
            "training_metrics": training_metrics,
            "trained_by": self.plugin_id
        }

        return {
            "model": trained_model,
            "metadata": model_metadata
        }


class ModelEvaluator(BasePlugin):
    """模型评估插件"""

    def execute(self, input_data: Any) -> dict:
        """评估模型性能"""
        # 处理多个依赖节点传入的数据
        if isinstance(input_data, dict):
            # 如果只有一个依赖节点的数据，直接使用
            if "model" in input_data and "X_test" in input_data and "y_test" in input_data:
                model = input_data["model"]
                X_test = input_data["X_test"]
                y_test = input_data["y_test"]
            # 如果有多个依赖节点的数据，需要从字典中提取
            else:
                # 从model_trainer节点获取模型
                model_data = input_data.get("model_trainer", {})
                model = model_data.get("model", {})

                # 从train_test_splitter节点获取测试数据
                splitter_data = input_data.get("train_test_splitter", {})
                X_test = splitter_data.get("X_test", [])
                y_test = splitter_data.get("y_test", [])
        else:
            # 如果不是字典格式，说明数据格式有问题
            raise ValueError("Invalid input data format for ModelEvaluator")

        evaluation_metrics = self.config.get("metrics", ["accuracy", "f1_score"])
        time.sleep(0.1)

        # 模拟模型评估
        test_metrics = {}
        for metric in evaluation_metrics:
            if metric == "accuracy":
                test_metrics["accuracy"] = random.uniform(0.80, 0.90)
            elif metric == "f1_score":
                test_metrics["f1_score"] = random.uniform(0.75, 0.85)
            elif metric == "precision":
                test_metrics["precision"] = random.uniform(0.78, 0.88)
            elif metric == "recall":
                test_metrics["recall"] = random.uniform(0.72, 0.82)

        evaluation_result = {
            "test_metrics": test_metrics,
            "test_sample_count": len(X_test),
            "model_type": model.get("model_type", "unknown") if isinstance(model, dict) else "unknown",
            "evaluated_by": self.plugin_id
        }

        return {
            "evaluation_result": evaluation_result
        }


class ModelRegistry(BasePlugin):
    """模型注册插件"""

    def execute(self, input_data: Any) -> dict:
        """注册模型"""
        evaluation_result = input_data["evaluation_result"]
        test_metrics = evaluation_result["test_metrics"]

        # 获取准确率
        accuracy = test_metrics.get("accuracy", 0)

        # 根据准确率决定是否注册模型
        if accuracy > 0.85:
            registry_type = "production"
            model_registered = True
            registration_reason = f"Model accuracy {accuracy:.3f} exceeds threshold 0.85"
        elif accuracy > 0.80:
            registry_type = "staging"
            model_registered = True
            registration_reason = f"Model accuracy {accuracy:.3f} meets staging criteria"
        else:
            registry_type = "rejected"
            model_registered = False
            registration_reason = f"Model accuracy {accuracy:.3f} below minimum threshold 0.80"

        time.sleep(0.05)  # 模拟注册时间

        registry_response = {
            "model_registered": model_registered,
            "registry_type": registry_type,
            "model_performance": test_metrics,
            "registration_reason": registration_reason,
            "registered_by": self.plugin_id
        }

        return {
            "registry_response": registry_response
        }